import numpy as np
import segyio
from scipy.signal.windows import triang
from scipy.signal import convolve2d as conv2


def read_sgy(file_path):
    """
    读取sgy文件，返回相关信息
    :param file_path:
    :return: {
            'cdpx': cdpx,  # cdpx
            'cdpy': cdpy,  # cdpx
            'wellx': wellx,  # wellx
            'welly': welly,  # welly
            'cdp': cdp,  # cdpx
            'fb': fb,  # 初至时间
            'inlin': inlin,  # 测线
            'crosslin': crosslin,  # 测线
            'hcf': hcf,  # 高频截止频率
            'dt': dt,  # 采样间隔
            'ntraces': ntraces,  # 地震道道数
            'nsamples': nsamples,  # 时间采样点数
            'tbasis': tbasis,  # 时间数组
            'data': data  # 地震数据
        }
    """
    with segyio.open(file_path, strict=False, ignore_geometry=True) as segy_file:
        data = np.stack([t.astype(float) for t in segy_file.trace])  # 将道数据浮点数，并堆叠成新的numpy数组
        cdp = np.array([t[segyio.TraceField.CDP_X] for t in segy_file.header])  # cdpx
        fb = np.array([t[segyio.TraceField.MuteTimeStart] for t in segy_file.header])  # 初至时间
        cdpx = np.array([t[segyio.TraceField.CDP_X] for t in segy_file.header])  # cdpx
        cdpy = np.array([t[segyio.TraceField.CDP_Y] for t in segy_file.header])  # cdpy
        wellx = np.array([t[segyio.TraceField.GroupX] for t in segy_file.header])  # wellx
        welly = np.array([t[segyio.TraceField.GroupY] for t in segy_file.header])  # welly
        inlin = np.array([t[segyio.TraceField.INLINE_3D] for t in segy_file.header])  # 测线号
        crosslin = np.array([t[segyio.TraceField.CROSSLINE_3D] for t in segy_file.header])  # 测线号
        hcf = np.array([t[segyio.TraceField.HighCutFrequency] for t in segy_file.header])  # 高频截止频率
        dt = segy_file.bin[segyio.BinField.Interval] / 1000  # 采样间隔

        data = np.asarray(data, dtype=float)
        ntraces = data.shape[-2]  # 地震道道数
        nsamples = int(data.shape[-1])  # 时间采样点数
        data = np.squeeze(data)  # 从数组的形状中移除长度为 1 的维度
        tbasis = np.arange(0, nsamples * dt, dt)  # 时间数组

        return {
            'cdpx': cdpx,  # cdpx
            'cdpy': cdpy,  # cdpx
            'wellx': wellx,  # wellx
            'welly': welly,  # welly
            'cdp': cdp,  # cdpx
            'fb': fb,  # 初至时间
            'inlin': inlin,  # 测线
            'crosslin': crosslin,  # 测线
            'hcf': hcf,  # 高频截止频率
            'dt': dt,  # 采样间隔
            'ntraces': ntraces,  # 地震道道数
            'nsamples': nsamples,  # 时间采样点数
            'tbasis': tbasis,  # 时间数组
            'data': data  # 地震数据
        }

def gain(file, agc):
    """
    对sgy数据进行增益处理
    :param file: sgy文件对象
    :param agc: 增益时窗默认为500
    :return: 返回处理后的道数据
    """
    data = file["data"]
    data = data.T  # 数组转置，交换数组的行和列
    dt = file["dt"]  # 采样间隔
    nt, nx = data.shape  # nt，时间采样点数，nx，道数
    dout = np.zeros(data.shape)  # 创建一个与data形状相同的零数组

    L = agc / dt + 1
    L = np.floor(L / 2)
    h = triang(2 * L + 1)  # triang用于生成三角形窗函数的函数
    shaped_h = h.reshape(len(h), 1)  #  将h的形状转换为(len(h), 1)

    # 对每一道进行处理
    for k in range(nx):
        aux = data[:, k]  # 第k道数据
        e = aux ** 2
        shaped_e = e.reshape(len(e), 1)

        rms = np.sqrt(conv2(shaped_e, shaped_h, "same"))
        epsi = 1e-10 * max(rms)
        op = rms / (rms ** 2 + epsi)
        op = op.reshape(len(op), )

        dout[:, k] = data[:, k] * op  # 将处理完的数据存放到dout中

    # 对每一道进行处理
    for k in range(nx):
        aux = dout[:, k]
        amax = np.sqrt(sum(aux ** 2) / nt)
        dout[:, k] = dout[:, k] / amax

    dout = dout.T

    return dout


